#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Tushare数据收集器
用于从Tushare获取各种金融数据，包括股票、基金、期货、财务报表等
"""

import os
import json
import yaml
import time
import logging
import pandas as pd
from datetime import datetime, timedelta
from pathlib import Path

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class TushareDataCollector:
    """Tushare数据收集器，用于从Tushare获取各种金融数据"""
    
    def __init__(self, config_path='config/config.yaml'):
        """
        初始化Tushare数据收集器
        
        Args:
            config_path: 配置文件路径
        """
        self.config_path = config_path
        self.config = self._load_config()
        self.tushare_config = self.config.get('tushare', {})
        self.stock_market_config = self.config.get('stock_market', {})
        
        # 创建保存目录
        self.output_path = self.config.get('system', {}).get('data_storage', {}).get('tushare_data_path', 'data/tushare_data/')
        os.makedirs(self.output_path, exist_ok=True)
        
        # 初始化Tushare
        self._init_tushare()
        
        logger.info("Tushare数据收集器初始化完成")
    
    def _load_config(self):
        """
        加载配置文件
        
        Returns:
            dict: 配置信息
        """
        try:
            with open(self.config_path, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)
            logger.info(f"成功加载配置文件: {self.config_path}")
            return config
        except Exception as e:
            logger.error(f"加载配置文件失败: {e}")
            return {}
    
    def _init_tushare(self):
        """
        初始化Tushare API
        
        Returns:
            bool: 是否成功
        """
        try:
            import tushare as ts
            
            # 从配置中获取API Token
            api_token = self.stock_market_config.get('api_token')
            if not api_token:
                logger.error("未找到Tushare API Token，请在配置文件中设置")
                self.tushare_api = None
                return False
            
            # 设置Token并获取API实例
            ts.set_token(api_token)
            self.tushare_api = ts.pro_api()
            
            logger.info("Tushare API初始化成功")
            return True
            
        except ImportError:
            logger.error("未安装Tushare，请使用pip install tushare安装")
            self.tushare_api = None
            return False
        except Exception as e:
            logger.error(f"初始化Tushare API失败: {e}")
            self.tushare_api = None
            return False
    
    def collect_from_stock_list(self, stock_list=None, data_types=None):
        """
        收集指定股票列表的各类数据
        
        Args:
            stock_list: 股票代码列表，如果为None则从配置中读取
            data_types: 要收集的数据类型列表，如果为None则从配置中读取
            
        Returns:
            bool: 是否成功
        """
        try:
            if self.tushare_api is None:
                logger.error("Tushare API未初始化，无法收集数据")
                return False
            
            if stock_list is None:
                # 从配置或分析结果中获取股票列表
                stock_list = self._get_stock_list_from_config()
            
            if not stock_list:
                logger.warning("股票列表为空，无法收集数据")
                return False
            
            if data_types is None:
                # 从配置中获取数据类型
                data_types = self.tushare_config.get('data_types', ['basic', 'financial', 'trading'])
            
            logger.info(f"准备收集 {len(stock_list)} 支股票的 {data_types} 数据")
            
            for stock_code in stock_list:
                logger.info(f"开始收集股票 {stock_code} 的数据")
                
                # 格式化股票代码
                ts_code = self._format_stock_code(stock_code)
                
                # 收集不同类型的数据
                stock_data = {}
                
                for data_type in data_types:
                    try:
                        if data_type == 'basic':
                            stock_data['basic'] = self._collect_basic_info(ts_code)
                        elif data_type == 'financial':
                            stock_data['financial'] = self._collect_financial_data(ts_code)
                        elif data_type == 'trading':
                            stock_data['trading'] = self._collect_trading_data(ts_code)
                        elif data_type == 'valuation':
                            stock_data['valuation'] = self._collect_valuation_data(ts_code)
                        elif data_type == 'forecast':
                            stock_data['forecast'] = self._collect_forecast_data(ts_code)
                        elif data_type == 'dividend':
                            stock_data['dividend'] = self._collect_dividend_data(ts_code)
                        elif data_type == 'major_holders':
                            stock_data['major_holders'] = self._collect_major_holders(ts_code)
                        elif data_type == 'industry':
                            stock_data['industry'] = self._collect_industry_data(ts_code)
                        else:
                            logger.warning(f"未知的数据类型: {data_type}")
                    except Exception as e:
                        logger.error(f"收集股票 {stock_code} 的 {data_type} 数据失败: {e}")
                
                # 保存单只股票的数据
                self._save_stock_data(stock_code, stock_data)
                
                # 避免请求过于频繁
                time.sleep(1)
            
            logger.info(f"成功收集 {len(stock_list)} 支股票的数据")
            return True
            
        except Exception as e:
            logger.error(f"收集股票数据失败: {e}")
            return False
    
    def collect_market_data(self, data_types=None):
        """
        收集市场整体数据，如指数、行业、宏观经济等
        
        Args:
            data_types: 要收集的数据类型列表，如果为None则从配置中读取
            
        Returns:
            bool: 是否成功
        """
        try:
            if self.tushare_api is None:
                logger.error("Tushare API未初始化，无法收集数据")
                return False
            
            if data_types is None:
                # 从配置中获取数据类型
                data_types = self.tushare_config.get('market_data_types', 
                                                  ['index', 'industry', 'macro', 'fund', 'futures'])
            
            logger.info(f"准备收集市场数据: {data_types}")
            
            market_data = {}
            
            for data_type in data_types:
                try:
                    if data_type == 'index':
                        market_data['index'] = self._collect_index_data()
                    elif data_type == 'industry':
                        market_data['industry'] = self._collect_industry_overview()
                    elif data_type == 'macro':
                        market_data['macro'] = self._collect_macro_data()
                    elif data_type == 'fund':
                        market_data['fund'] = self._collect_fund_data()
                    elif data_type == 'futures':
                        market_data['futures'] = self._collect_futures_data()
                    elif data_type == 'crypto':
                        market_data['crypto'] = self._collect_crypto_data()
                    else:
                        logger.warning(f"未知的市场数据类型: {data_type}")
                except Exception as e:
                    logger.error(f"收集 {data_type} 市场数据失败: {e}")
            
            # 保存市场数据
            self._save_market_data(market_data)
            
            logger.info(f"成功收集市场数据")
            return True
            
        except Exception as e:
            logger.error(f"收集市场数据失败: {e}")
            return False
    
    def _get_stock_list_from_config(self):
        """
        从配置或分析结果中获取股票列表
        
        Returns:
            list: 股票代码列表
        """
        try:
            # 首先尝试从配置中获取
            stock_list = self.tushare_config.get('stock_list', [])
            
            # 如果配置中没有，则从最新的分析结果中获取
            if not stock_list:
                results_path = self.config.get('system', {}).get('data_storage', {}).get('results_path', 'data/results/')
                date = datetime.now().strftime('%Y%m%d')
                
                # 尝试获取今天的数据，如果没有则获取昨天的
                analysis_file = os.path.join(results_path, f'analysis_results_{date}.json')
                if not os.path.exists(analysis_file):
                    date = (datetime.now() - timedelta(days=1)).strftime('%Y%m%d')
                    analysis_file = os.path.join(results_path, f'analysis_results_{date}.json')
                
                if os.path.exists(analysis_file):
                    with open(analysis_file, 'r', encoding='utf-8') as f:
                        analysis_data = json.load(f)
                    stock_list = list(analysis_data.keys())
            
            # 限制数量，避免一次抓取过多
            max_stocks = self.tushare_config.get('max_stocks_per_collection', 50)
            if len(stock_list) > max_stocks:
                stock_list = stock_list[:max_stocks]
                logger.info(f"限制股票数量为 {max_stocks}")
            
            return stock_list
            
        except Exception as e:
            logger.error(f"获取股票列表失败: {e}")
            return []
    
    def _format_stock_code(self, stock_code):
        """
        格式化股票代码为Tushare格式
        
        Args:
            stock_code: 原始股票代码
            
        Returns:
            str: 格式化后的股票代码
        """
        # Tushare格式: 000001.SZ, 600000.SH, 300000.SZ, 688001.SH
        if '.' in stock_code:
            return stock_code
        
        # 深市股票(000, 001, 002, 003, 300, 301开头)
        if stock_code.startswith(('000', '001', '002', '003', '300', '301')):
            return f"{stock_code}.SZ"
        # 上海科创板(688开头)
        elif stock_code.startswith('688'):
            return f"{stock_code}.SH"
        # 上海主板(600, 601, 603, 605开头)
        elif stock_code.startswith(('600', '601', '603', '605')):
            return f"{stock_code}.SH"
        # 北交所(8开头)
        elif stock_code.startswith('8'):
            return f"{stock_code}.BJ"
        # 港股
        elif stock_code.startswith('00') and len(stock_code) <= 5:
            return f"{stock_code}.HK"
        else:
            return stock_code
    
    def _collect_basic_info(self, ts_code):
        """
        收集股票基本信息
        
        Args:
            ts_code: Tushare格式的股票代码
            
        Returns:
            dict: 股票基本信息
        """
        try:
            # 获取股票基本信息
            df = self.tushare_api.stock_company(ts_code=ts_code)
            if df.empty:
                logger.warning(f"未找到股票 {ts_code} 的基本信息")
                return {}
            
            # 获取股票详细信息
            df_basic = self.tushare_api.stock_basic(ts_code=ts_code, fields='ts_code,symbol,name,area,industry,fullname,enname,market,exchange,curr_type,list_status,list_date,delist_date,is_hs')
            
            # 合并信息
            result = {}
            if not df.empty:
                result.update(df.iloc[0].to_dict())
            if not df_basic.empty:
                result.update(df_basic.iloc[0].to_dict())
            
            return result
            
        except Exception as e:
            logger.error(f"收集股票 {ts_code} 的基本信息失败: {e}")
            return {}
    
    def _collect_financial_data(self, ts_code):
        """
        收集股票财务数据
        
        Args:
            ts_code: Tushare格式的股票代码
            
        Returns:
            dict: 股票财务数据
        """
        try:
            # 获取最近的财务数据
            # 资产负债表
            df_balance = self.tushare_api.balancesheet(ts_code=ts_code, period='')
            # 利润表
            df_income = self.tushare_api.income(ts_code=ts_code, period='')
            # 现金流量表
            df_cashflow = self.tushare_api.cashflow(ts_code=ts_code, period='')
            # 主要财务指标
            df_indicator = self.tushare_api.fina_indicator(ts_code=ts_code, period='')
            
            result = {
                'balance_sheet': [] if df_balance.empty else df_balance.head(4).to_dict('records'),
                'income': [] if df_income.empty else df_income.head(4).to_dict('records'),
                'cashflow': [] if df_cashflow.empty else df_cashflow.head(4).to_dict('records'),
                'indicator': [] if df_indicator.empty else df_indicator.head(4).to_dict('records')
            }
            
            return result
            
        except Exception as e:
            logger.error(f"收集股票 {ts_code} 的财务数据失败: {e}")
            return {}
    
    def _collect_trading_data(self, ts_code):
        """
        收集股票交易数据
        
        Args:
            ts_code: Tushare格式的股票代码
            
        Returns:
            dict: 股票交易数据
        """
        try:
            # 获取近期日线行情
            end_date = datetime.now().strftime('%Y%m%d')
            start_date = (datetime.now() - timedelta(days=60)).strftime('%Y%m%d')
            
            df_daily = self.tushare_api.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
            
            # 获取近期指标数据
            df_daily_basic = self.tushare_api.daily_basic(ts_code=ts_code, start_date=start_date, end_date=end_date)
            
            result = {
                'daily': [] if df_daily.empty else df_daily.to_dict('records'),
                'daily_basic': [] if df_daily_basic.empty else df_daily_basic.to_dict('records')
            }
            
            return result
            
        except Exception as e:
            logger.error(f"收集股票 {ts_code} 的交易数据失败: {e}")
            return {}
    
    def _collect_valuation_data(self, ts_code):
        """
        收集股票估值数据
        
        Args:
            ts_code: Tushare格式的股票代码
            
        Returns:
            dict: 股票估值数据
        """
        try:
            # 获取市盈率和市净率数据
            end_date = datetime.now().strftime('%Y%m%d')
            start_date = (datetime.now() - timedelta(days=60)).strftime('%Y%m%d')
            
            df_daily_basic = self.tushare_api.daily_basic(ts_code=ts_code, start_date=start_date, end_date=end_date,
                                                       fields='ts_code,trade_date,close,turnover_rate,turnover_rate_f,volume_ratio,pe,pe_ttm,pb,ps,ps_ttm,dv_ratio,dv_ttm,total_mv,circ_mv')
            
            result = {
                'valuation': [] if df_daily_basic.empty else df_daily_basic.to_dict('records')
            }
            
            return result
            
        except Exception as e:
            logger.error(f"收集股票 {ts_code} 的估值数据失败: {e}")
            return {}
    
    def _collect_forecast_data(self, ts_code):
        """
        收集股票业绩预测数据
        
        Args:
            ts_code: Tushare格式的股票代码
            
        Returns:
            dict: 股票业绩预测数据
        """
        try:
            # 获取业绩预告数据
            df_forecast = self.tushare_api.forecast(ts_code=ts_code)
            
            # 获取业绩快报数据
            df_express = self.tushare_api.express(ts_code=ts_code)
            
            result = {
                'forecast': [] if df_forecast.empty else df_forecast.head(4).to_dict('records'),
                'express': [] if df_express.empty else df_express.head(4).to_dict('records')
            }
            
            return result
            
        except Exception as e:
            logger.error(f"收集股票 {ts_code} 的业绩预测数据失败: {e}")
            return {}
    
    def _collect_dividend_data(self, ts_code):
        """
        收集股票分红数据
        
        Args:
            ts_code: Tushare格式的股票代码
            
        Returns:
            dict: 股票分红数据
        """
        try:
            # 获取分红数据
            df_dividend = self.tushare_api.dividend(ts_code=ts_code)
            
            result = {
                'dividend': [] if df_dividend.empty else df_dividend.head(10).to_dict('records')
            }
            
            return result
            
        except Exception as e:
            logger.error(f"收集股票 {ts_code} 的分红数据失败: {e}")
            return {}
    
    def _collect_major_holders(self, ts_code):
        """
        收集股票主要股东数据
        
        Args:
            ts_code: Tushare格式的股票代码
            
        Returns:
            dict: 股票主要股东数据
        """
        try:
            # 获取十大股东数据
            df_top10_holders = self.tushare_api.top10_holders(ts_code=ts_code)
            
            # 获取十大流通股东数据
            df_top10_floatholders = self.tushare_api.top10_floatholders(ts_code=ts_code)
            
            # 获取机构持股数据
            df_institutional_investors = self.tushare_api.query('stk_holdertrade', ts_code=ts_code)
            
            result = {
                'top10_holders': [] if df_top10_holders.empty else df_top10_holders.head(20).to_dict('records'),
                'top10_floatholders': [] if df_top10_floatholders.empty else df_top10_floatholders.head(20).to_dict('records'),
                'institutional_investors': [] if df_institutional_investors.empty else df_institutional_investors.head(20).to_dict('records')
            }
            
            return result
            
        except Exception as e:
            logger.error(f"收集股票 {ts_code} 的主要股东数据失败: {e}")
            return {}
    
    def _collect_industry_data(self, ts_code):
        """
        收集股票所在行业数据
        
        Args:
            ts_code: Tushare格式的股票代码
            
        Returns:
            dict: 股票所在行业数据
        """
        try:
            # 获取基本信息以确定行业
            df_basic = self.tushare_api.stock_basic(ts_code=ts_code, fields='industry')
            
            if df_basic.empty:
                logger.warning(f"未找到股票 {ts_code} 的行业信息")
                return {}
            
            industry = df_basic.iloc[0]['industry']
            
            # 获取行业内其他股票
            df_industry_stocks = self.tushare_api.stock_basic(industry=industry, fields='ts_code,symbol,name,area,industry,list_date')
            
            result = {
                'industry': industry,
                'industry_stocks': [] if df_industry_stocks.empty else df_industry_stocks.head(20).to_dict('records')
            }
            
            return result
            
        except Exception as e:
            logger.error(f"收集股票 {ts_code} 的行业数据失败: {e}")
            return {}
    
    def _collect_index_data(self):
        """
        收集指数数据
        
        Returns:
            dict: 指数数据
        """
        try:
            # 主要指数列表
            index_codes = ['000001.SH', '399001.SZ', '000300.SH', '000905.SH', '399006.SZ']
            
            end_date = datetime.now().strftime('%Y%m%d')
            start_date = (datetime.now() - timedelta(days=30)).strftime('%Y%m%d')
            
            result = {}
            
            for index_code in index_codes:
                try:
                    # 获取指数日线数据
                    df_index = self.tushare_api.index_daily(ts_code=index_code, start_date=start_date, end_date=end_date)
                    
                    # 指数名称
                    index_name = {
                        '000001.SH': '上证指数',
                        '399001.SZ': '深证成指',
                        '000300.SH': '沪深300',
                        '000905.SH': '中证500',
                        '399006.SZ': '创业板指'
                    }.get(index_code, index_code)
                    
                    result[index_code] = {
                        'name': index_name,
                        'data': [] if df_index.empty else df_index.to_dict('records')
                    }
                    
                except Exception as e:
                    logger.error(f"收集指数 {index_code} 数据失败: {e}")
            
            return result
            
        except Exception as e:
            logger.error(f"收集指数数据失败: {e}")
            return {}
    
    def _collect_industry_overview(self):
        """
        收集行业概览数据
        
        Returns:
            dict: 行业概览数据
        """
        try:
            # 获取行业分类
            df_industry = self.tushare_api.stock_basic(fields='industry')
            industry_list = df_industry['industry'].unique().tolist()
            
            result = {}
            
            for industry in industry_list:
                if not industry or pd.isna(industry):
                    continue
                    
                try:
                    # 获取行业内股票
                    df_stocks = self.tushare_api.stock_basic(industry=industry, fields='ts_code,name,area,industry,list_date,market')
                    
                    result[industry] = {
                        'name': industry,
                        'stock_count': 0 if df_stocks.empty else len(df_stocks),
                        'stocks': [] if df_stocks.empty else df_stocks.head(10).to_dict('records')
                    }
                    
                except Exception as e:
                    logger.error(f"收集行业 {industry} 概览数据失败: {e}")
            
            return result
            
        except Exception as e:
            logger.error(f"收集行业概览数据失败: {e}")
            return {}
    
    def _collect_macro_data(self):
        """
        收集宏观经济数据
        
        Returns:
            dict: 宏观经济数据
        """
        try:
            # 获取GDP数据
            df_gdp = self.tushare_api.query('monthly_data', data_name='CN/GDP')
            
            # 获取CPI数据
            df_cpi = self.tushare_api.query('monthly_data', data_name='CN/CPI')
            
            # 获取PPI数据
            df_ppi = self.tushare_api.query('monthly_data', data_name='CN/PPI')
            
            # 获取货币供应量数据
            df_m2 = self.tushare_api.query('monthly_data', data_name='CN/M2')
            
            result = {
                'gdp': [] if df_gdp.empty else df_gdp.head(10).to_dict('records'),
                'cpi': [] if df_cpi.empty else df_cpi.head(10).to_dict('records'),
                'ppi': [] if df_ppi.empty else df_ppi.head(10).to_dict('records'),
                'm2': [] if df_m2.empty else df_m2.head(10).to_dict('records')
            }
            
            return result
            
        except Exception as e:
            logger.error(f"收集宏观经济数据失败: {e}")
            return {}
    
    def _collect_fund_data(self):
        """
        收集基金数据
        
        Returns:
            dict: 基金数据
        """
        try:
            # 获取基金基本信息
            df_fund_basic = self.tushare_api.fund_basic(market='E')
            
            # 获取基金经理信息
            df_fund_manager = self.tushare_api.fund_manager()
            
            # 获取基金持仓信息(季度更新)
            df_fund_portfolio = self.tushare_api.fund_portfolio()
            
            # 获取基金净值信息
            today = datetime.now().strftime('%Y%m%d')
            df_fund_nav = self.tushare_api.fund_nav(end_date=today)
            
            result = {
                'fund_basic': [] if df_fund_basic.empty else df_fund_basic.head(50).to_dict('records'),
                'fund_manager': [] if df_fund_manager.empty else df_fund_manager.head(50).to_dict('records'),
                'fund_portfolio': [] if df_fund_portfolio.empty else df_fund_portfolio.head(50).to_dict('records'),
                'fund_nav': [] if df_fund_nav.empty else df_fund_nav.head(50).to_dict('records')
            }
            
            return result
            
        except Exception as e:
            logger.error(f"收集基金数据失败: {e}")
            return {}
    
    def _collect_futures_data(self):
        """
        收集期货数据
        
        Returns:
            dict: 期货数据
        """
        try:
            # 获取期货合约信息
            df_future_basic = self.tushare_api.fut_basic()
            
            # 获取主要期货合约数据
            result = {
                'future_basic': [] if df_future_basic.empty else df_future_basic.head(50).to_dict('records'),
                'future_daily': {}
            }
            
            # 主要期货合约
            main_futures = ['CU', 'AL', 'ZN', 'PB', 'NI', 'SN', 'AU', 'AG', 'RB', 'HC', 'FU', 'BU', 'RU', 'A', 'M', 'Y', 'P', 'C', 'CS', 'JD']
            
            end_date = datetime.now().strftime('%Y%m%d')
            start_date = (datetime.now() - timedelta(days=30)).strftime('%Y%m%d')
            
            for future_code in main_futures[:5]:  # 限制数量
                try:
                    # 获取期货日线数据
                    df_future_daily = self.tushare_api.fut_daily(ts_code=f'{future_code}.SHF', start_date=start_date, end_date=end_date)
                    
                    result['future_daily'][future_code] = [] if df_future_daily.empty else df_future_daily.to_dict('records')
                    
                except Exception as e:
                    logger.error(f"收集期货 {future_code} 数据失败: {e}")
            
            return result
            
        except Exception as e:
            logger.error(f"收集期货数据失败: {e}")
            return {}
    
    def _collect_crypto_data(self):
        """
        收集数字货币数据
        
        Returns:
            dict: 数字货币数据
        """
        try:
            # 由于Tushare对数字货币支持有限，这里使用模拟数据
            result = {
                'crypto_daily': {
                    'BTC': {
                        'name': '比特币',
                        'data': [
                            {'date': '2023-05-01', 'close': 29000, 'volume': 30000},
                            {'date': '2023-05-02', 'close': 28500, 'volume': 32000},
                            {'date': '2023-05-03', 'close': 28800, 'volume': 28000}
                        ]
                    },
                    'ETH': {
                        'name': '以太坊',
                        'data': [
                            {'date': '2023-05-01', 'close': 1800, 'volume': 15000},
                            {'date': '2023-05-02', 'close': 1750, 'volume': 16000},
                            {'date': '2023-05-03', 'close': 1820, 'volume': 14000}
                        ]
                    }
                }
            }
            
            return result
            
        except Exception as e:
            logger.error(f"收集数字货币数据失败: {e}")
            return {}
    
    def _save_stock_data(self, stock_code, stock_data):
        """
        保存单只股票的数据
        
        Args:
            stock_code: 股票代码
            stock_data: 股票数据字典
            
        Returns:
            bool: 是否成功
        """
        try:
            # 创建股票目录
            stock_dir = os.path.join(self.output_path, 'stocks', stock_code)
            os.makedirs(stock_dir, exist_ok=True)
            
            date = datetime.now().strftime('%Y%m%d')
            
            # 保存不同类型的数据到不同文件
            for data_type, data in stock_data.items():
                output_file = os.path.join(stock_dir, f'{data_type}_{date}.json')
                
                with open(output_file, 'w', encoding='utf-8') as f:
                    json.dump(data, f, ensure_ascii=False, indent=2)
                
                logger.info(f"成功保存股票 {stock_code} 的 {data_type} 数据到: {output_file}")
            
            return True
            
        except Exception as e:
            logger.error(f"保存股票 {stock_code} 数据失败: {e}")
            return False
    
    def _save_market_data(self, market_data):
        """
        保存市场数据
        
        Args:
            market_data: 市场数据字典
            
        Returns:
            bool: 是否成功
        """
        try:
            # 创建市场数据目录
            market_dir = os.path.join(self.output_path, 'market')
            os.makedirs(market_dir, exist_ok=True)
            
            date = datetime.now().strftime('%Y%m%d')
            
            # 保存不同类型的数据到不同文件
            for data_type, data in market_data.items():
                output_file = os.path.join(market_dir, f'{data_type}_{date}.json')
                
                with open(output_file, 'w', encoding='utf-8') as f:
                    json.dump(data, f, ensure_ascii=False, indent=2)
                
                logger.info(f"成功保存 {data_type} 市场数据到: {output_file}")
            
            # 保存汇总数据
            summary = {
                'date': date,
                'data_types': list(market_data.keys()),
                'success': True
            }
            
            summary_file = os.path.join(market_dir, f'summary_{date}.json')
            with open(summary_file, 'w', encoding='utf-8') as f:
                json.dump(summary, f, ensure_ascii=False, indent=2)
            
            return True
            
        except Exception as e:
            logger.error(f"保存市场数据失败: {e}")
            return False

def main():
    """主函数"""
    collector = TushareDataCollector()
    collector.collect_from_stock_list()
    collector.collect_market_data()

if __name__ == '__main__':
    main() 